package com.wfh.aiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.wfh.aiagent.model.Conversation;
import com.wfh.aiagent.service.ConversationService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.stereotype.Component;

import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

/**
 * 基于数据库的聊天记忆实现，使用 Kryo 序列化消息列表。
 *
 * @Author FengHuan Wang
 * @Date 2025/5/25 13:24
 * @Version 1.0
 */
@Component
@Slf4j
public class DbBasedMemory implements ChatMemory {

    @Resource
    private ConversationService conversationService;

    private static final Kryo kryo = new Kryo();

    static {
        kryo.setRegistrationRequired(false);
        kryo.register(Conversation.class);
        kryo.register(ArrayList.class);
        kryo.register(MessageType.class);
        kryo.register(HashMap.class);
        kryo.register(org.springframework.ai.chat.messages.UserMessage.class);
        kryo.register(org.springframework.ai.chat.messages.AssistantMessage.class);
        kryo.register(org.springframework.ai.chat.messages.SystemMessage.class);
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    /**
     * 获取指定会话的全部消息（必须实现，来自 ChatMemory 接口）
     */
    public List<Message> get(String conversationId) {
        return getOrCreateConversation(conversationId);
    }

    /**
     * 添加消息到指定会话（必须实现，来自 ChatMemory 接口）
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        if (messages == null || messages.isEmpty()) {
            return;
        }
        List<Message> existingMessages = getOrCreateConversation(conversationId);
        existingMessages.addAll(messages);
        saveConversation(conversationId, existingMessages);
    }

    /**
     * 辅助方法：添加单条消息（非接口方法，不加 @Override）
     */
    @Override
    public void add(String conversationId, Message message) {
        if (message != null) {
            add(conversationId, List.of(message));
        }
    }

    /**
     * 获取指定会话的最近 N 条消息（自定义扩展方法）
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> messageList = getOrCreateConversation(conversationId);
        int fromIndex = Math.max(0, messageList.size() - lastN);
        return new ArrayList<>(messageList.subList(fromIndex, messageList.size()));
    }

    /**
     * 清除指定会话的所有消息（必须实现，来自 ChatMemory 接口）
     */
    @Override
    public void clear(String conversationId) {
        try {
            Conversation conversation = conversationService.getById(conversationId);
            if (conversation != null) {
                conversationService.getBaseMapper().deleteById(conversation.getId());
            }
        } catch (Exception e) {
            log.error("清除会话 {} 失败", conversationId, e);
        }
    }

    /**
     * 从数据库获取或创建一个会话的消息列表
     */
    private List<Message> getOrCreateConversation(String conversationId) {
        try {
            Conversation conversation = conversationService.getById(conversationId);
            if (conversation != null && conversation.getMessage() != null) {
                String message = conversation.getMessage();
                byte[] bytes = message.getBytes(StandardCharsets.UTF_8);
                try (Input input = new Input(bytes)) {
                    @SuppressWarnings("unchecked")
                    List<Message> messages = kryo.readObject(input, ArrayList.class);
                    return messages != null ? messages : new ArrayList<>();
                } catch (Exception e) {
                    log.error("反序列化会话 {} 失败", conversationId, e);
                }
            }
        } catch (Exception e) {
            log.error("从数据库加载会话 {} 出错", conversationId, e);
        }
        return new ArrayList<>();
    }

    /**
     * 将消息列表保存到数据库
     */
    private void saveConversation(String conversationId, List<Message> messages) {
        try (Output output = new Output(1024, -1)) {
            kryo.writeObject(output, messages);
            byte[] bytes = output.toBytes();
            Conversation conversation = new Conversation();
            conversation.setId(Integer.valueOf(conversationId));
            conversation.setMessage(new String(bytes, StandardCharsets.UTF_8));

            if (conversationService.getById(conversationId) != null) {
                conversationService.updateById(conversation);
            } else {
                conversationService.save(conversation);
            }
        } catch (Exception e) {
            log.error("保存会话 {} 失败", conversationId, e);
        }
    }
}